import requests
from wxauto import *
import time
import random
import os
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import re


def get_chat_history(contact_name):
    """
    获取指定联系人或群聊的聊天记录并保存到本地文件。

    参数:
    contact_name (str): 联系人或群聊的名称

    返回:
    str: 保存聊天记录的文件路径
    """
    # 创建 WeChat 对象，用于操作微信
    wx = WeChat()
    # 切换到指定联系人或群聊的聊天窗口
    wx.ChatWith(contact_name)
    # 等待2秒，确保聊天窗口加载完成
    time.sleep(2)

    all_msg = []
    # 获取当前聊天窗口的所有消息
    messages = wx.GetAllMessage()
    for message in messages:
        # 将消息按格式组合后添加到 all_msg 列表中
        all_msg.append(f"{message[0]}:{message[1]}")
    # 获取当前目录
    path = os.getcwd()
    # 拼接保存聊天记录的文件路径
    file_path = os.path.join(path, f"{contact_name}所有聊天记录.txt")
    # 以写入模式打开文件，将聊天记录写入文件
    with open(file_path, 'w', encoding='utf-8') as f:
        for item in all_msg:
            f.write(item + '\n')

    return file_path


# 用于读取某个人的消息来模仿其风格
def read_style_examples(file_path):
    """
    从指定文件中读取聊天记录，用于模仿某人的聊天风格。

    参数:
    file_path (str): 聊天记录文件的路径

    返回:
    list: 包含聊天记录的列表，每个元素是一个字典，包含 "role" 和 "content" 键
    """
    examples = []
    # 以只读模式打开文件
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            # 去除每行首尾的空白字符
            line = line.strip()
            if line:
                # 将每行聊天记录添加到 examples 列表中
                examples.append({"role": "user", "content": line})
    return examples


name = "不可一日无笑"
path = os.getcwd()
file = name + "所有聊天记录.txt"
now_file = path + "\\" + file
if not os.path.exists(now_file):
    # 如果聊天记录文件不存在，则获取聊天记录
    file_path = get_chat_history(name)
    print("聊天记录已经获取成功")
else:
    file_path = now_file
    print("所有聊天记录已经获取成功")
# 初始化历史对话列表，添加目标人物风格示例消息
style_examples = read_style_examples(file_path)
instruction = {"role": "user", "content": f"请模仿{name}的语气进行回复"}
history_messages = [instruction] + style_examples.copy()


def get_spark_reply(last_message):
    """
    使用星火大模型获取回复，并更新历史对话列表。

    参数:
    last_message (str): 最新的用户消息

    返回:
    str: 星火大模型的回复
    """
    global history_messages
    # 将新消息添加到历史对话列表
    history_messages.append({"role": "user", "content": last_message})
    url = "https://spark-api-open.xf-yun.com/v1/chat/completions"
    data = {
        "model": "4.0Ultra",
        "messages": history_messages,
        "temperature": 0.8,
        "top_p": 1,
        "max_tokens": 200,
        "frequency_penalty": 0.2,
        "presence_penalty": 0.2
    }
    # 从环境变量中获取 API 密钥
    api_password = "Bearer kDGPfkUvmLgoJFYSMmlN:KObNnGCjgYoSCCkKXzFH"

    header = {
        "Content-Type": "application/json",
        "Authorization": f"{api_password}"
    }

    try:
        # 发送 POST 请求获取回复
        response = requests.post(url, headers=header, json=data)
        # 获取模型回复并添加到历史对话列表
        result = response.json()
        history_messages.append({"role": "assistant", "content": result["choices"][0]["message"]["content"]})
        return result["choices"][0]["message"]["content"]
    except requests.RequestException as e:
        print(f"请求出错: {e}")
        error_replies = ["我这里网不好", "得,网不好"]
        return random.choice(error_replies)

    except (KeyError, IndexError) as e:
        print(f"解析响应出错: {e}")
        error_replies = ["不理解，听不懂你在说什么", "sorry啦，再说一遍呗"]
        return random.choice(error_replies)


def preprocess_message(message, tokenizer):
    """
    对输入的消息进行预处理，将其转换为模型可以接受的输入格式。

    参数:
    message (str): 输入的消息
    tokenizer: 分词器对象

    返回:
    torch.Tensor: 处理后的输入张量
    """
    inputs = tokenizer(message, truncation=True, padding='max_length', return_tensors='pt')
    return inputs


def predict_sentiment(message, tokenizer, model):
    """
    预测输入消息的情感倾向。

    参数:
    message (str): 输入的消息
    tokenizer: 分词器对象
    model: 情感分类模型

    返回:
    int: 预测的情感标签，0 表示负面，1 表示中性，2 表示正面
    """
    inputs = preprocess_message(message, tokenizer)
    with torch.no_grad():
        # 进行前向传播，获取模型输出
        outputs = model(**inputs)
        logits = outputs.logits
        # 获取预测的情感标签
        _, predicted = torch.max(logits, dim=1)
        return predicted.item()


def main():
    """
    主函数，用于启动微信自动回复程序。
    """
    model_path = os.path.join(path, f"results/checkpoint-264")
    if not os.path.exists(model_path):
        print(f"指定的模型路径 {model_path} 不存在，请检查。")
        return
    # 加载分词器
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    # 加载情感分类模型
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    # 将模型设置为评估模式
    model.eval()
    # 创建 WeChat 对象，用于操作微信
    wx = WeChat()
    while True:
        # 获取最新消息
        last_message = wx.GetAllNewMessage()
        # 获取当前聊天对象的名字
        contact_name = wx.CurrentChat()
        if last_message:
            for message in last_message[contact_name]:
                sender = message[0]
                msg = message[1]
                if last_message:
                    # 判断是谁回复的，如果是自己或者SYS（system）则忽略，否则会回复自己发的消息
                    if sender == contact_name:
                        print(f"收到消息: {msg}")
                        if re.match(r'^[。]+$', msg):
                            ai_reply = "脑子有泡吗，一直冒泡"
                            print(f"AI 回复: {ai_reply}")
                            # 向联系人发送回复
                            wx.SendMsg(str(ai_reply))
                            continue
                        # 预测消息的情感倾向
                        sentiment = predict_sentiment(msg, tokenizer, model)
                        if sentiment == 0:
                            print("消息情感倾向为负面")
                            # 可以根据负面情感调整回复策略
                        elif sentiment == 1:
                            print("消息情感倾向为中性")
                        else:
                            print("消息情感倾向为正面")
                        # 调用星火大模型获取回复
                        ai_reply = get_spark_reply(msg)
                        print(f"AI 回复: {ai_reply}")
                        # 向联系人发送回复
                        wx.SendMsg(str(ai_reply))
                # 休眠1秒，避免频繁操作
                time.sleep(1.5)


if __name__ == "__main__":
    # 创建 WeChat 对象
    wx = WeChat()
    main()
